#鲸鱼优化算法
import numpy as np
import random
import math
import copy
import time

class WOA():
    def __init__(self,n_dim=None, pop_size=20, max_iter=50, lb=-512, ub=512,obj_func=None):
        self.dim = n_dim
        self.pop = pop_size
        self.ub = ub
        self.lb = lb
        self.Max_iter = max_iter
        self.fun = obj_func
        self.X = np.zeros([pop_size,n_dim])
        for i in range(pop_size):
            for j in range(n_dim):
                self.X[i,j] = (ub[j] - lb[j]) * np.random.random() + lb[j]

    def BorderCheck(self):
        #边界检查
        pop = self.pop
        dim = self.dim
        lb = self.lb
        ub = self.ub
        for i in range(pop):
            for j in range(dim):
                if self.X[i,j] > ub[j]:
                    self.X[i, j] = ub[j]
                elif self.X[i,j] < lb[j]:
                    self.X[i, j] = lb[j]


    '''计算适应度函数'''
    def CaculateFitness(self):
        pop = self.X.shape[0]
        fitness = np.zeros([pop, 1])
        for i in range(pop):
            fitness[i] = self.fun(self.X[i, :])
        return fitness

    '''对适应度排序'''
    def SortFitness(self,Fit):
        fitness = np.sort(Fit, axis=0)
        index = np.argsort(Fit, axis=0)
        return fitness, index

    '''根据适应度对位置进行排序'''

    def SortPosition(self, index):
        Xnew = np.zeros(self.X.shape)
        for i in range(self.X.shape[0]):
            Xnew[i, :] = self.X[index[i], :]
        return Xnew

    '''鲸鱼优化算法本体'''

    def Tent_WOA(self):
        pop = self.pop
        dim = self.dim
        ub = self.ub
        lb = self.lb
        fun = self.fun
        maxIter = self.Max_iter
        fitness = self.CaculateFitness()
        fitness,sortIndex = self.SortFitness(fitness)
        self.X = self.SortPosition(sortIndex)
        GbestScore = copy.copy(fitness[0])
        GbestPositon = np.zeros([1,dim])
        GbestPositon[0,:] = copy.copy(self.X[0,:])
        Curve = np.zeros([maxIter,1])
        avg_fit = []  # 每次迭代的所有个体适应度平均值
        avg_ts = []  # 平均方差
        time_start = time.time()  # 记录迭代寻优开始时间
        for t in range(maxIter):
            Leader = self.X[0,:]
            a = 2-t*(2/maxIter)
            for i in range(pop):
                r1 = random.random()
                r2 = random.random()

                A = 2*a*r1 - a
                C = 2*r2
                b = 1
                l = 2*random.random() - 1

                for j in range(dim):
                    p = random.random()
                    if p < 0.5:
                        if np.abs(A) >= 1:   #寻找猎物
                            rand_leader_index = min(int(np.floor(pop*random.random() + 1)),pop - 1)   #随机选择一个个体
                            X_rand = self.X[rand_leader_index,:]
                            D_X_rand = np.abs(C*X_rand[j] - self.X[i,j])
                            self.X[i,j] = X_rand[j] - A*D_X_rand
                        elif np.abs(A) < 1:   #包围猎物
                            D_Leader = np.abs(C*Leader[j] - self.X[i,j])
                            self.X[i,j] = Leader[j] - A*D_Leader
                    elif p>= 0.5:   #气泡网攻击
                        distance2Leader = np.abs(Leader[j] - self.X[i,j])
                        self.X[i,j] = distance2Leader*np.exp(b*l)*np.cos(l*2*math.pi)+Leader[j]

            self.BorderCheck()
            fitness = self.CaculateFitness()
            avg_fit.append(np.mean(fitness))
            avg_ts.append(np.var(fitness))
            fitness, sortIndex = self.SortFitness(fitness)
            self.X = self.SortPosition(sortIndex)
            if fitness[0] <= GbestScore:
                GbestScore = copy.copy(fitness[0])
                GbestPositon[0,:] = copy.copy(self.X[0,:])
            Curve[t] = GbestScore
        time_end = time.time()  # 记录迭代结束时间
        print(f'WOA共花费 {time_end - time_start} 秒')

        return Curve,avg_fit,avg_ts







